{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64ec2f0f",
   "metadata": {},
   "source": [
    "# Lesson 3: Building an Agent Reasoning Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0d7f1cf",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b07baa43-7a51-4c39-91cc-aa0d9619b69f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from helper import get_openai_api_key\n",
    "OPENAI_API_KEY = get_openai_api_key()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcfa86a3-c7f2-41fa-b8b6-5617659ec36a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d3af4bb",
   "metadata": {},
   "source": [
    "## Load the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d8bfb34",
   "metadata": {
    "tags": []
   },
   "source": [
    "To download this paper, below is the needed code:\n",
    "\n",
    "#!wget \"https://openreview.net/pdf?id=VtmBAGCN7o\" -O metagpt.pdf\n",
    "\n",
    "**Note**: The pdf file is included with this lesson. To access it, go to the `File` menu and select`Open...`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb741560",
   "metadata": {},
   "source": [
    "## Setup the Query Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77464fb2-5ace-4839-9032-a020df8d4259",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import get_doc_tools\n",
    "\n",
    "vector_tool, summary_tool = get_doc_tools(\"metagpt.pdf\", \"metagpt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40aae3fc",
   "metadata": {},
   "source": [
    "## Setup Function Calling Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff4f5199-d02c-47b0-a9ab-cf72c8a506a3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9365d78d-8e9f-4f22-8d57-35a4c6aa6baf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.core.agent import FunctionCallingAgentWorker\n",
    "from llama_index.core.agent import AgentRunner\n",
    "\n",
    "agent_worker = FunctionCallingAgentWorker.from_tools(\n",
    "    [vector_tool, summary_tool], \n",
    "    llm=llm, \n",
    "    verbose=True\n",
    ")\n",
    "agent = AgentRunner(agent_worker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a9535d7-0baf-4905-ad16-5fb903d33b85",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = agent.query(\n",
    "    \"Tell me about the agent roles in MetaGPT, \"\n",
    "    \"and then how they communicate with each other.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcf74ec4-559f-4284-9ed0-817d26951c54",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(response.source_nodes[0].get_content(metadata_mode=\"all\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b28c184-0b65-4e38-808e-d91a285aaefe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = agent.chat(\n",
    "    \"Tell me about the evaluation datasets used.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9586cef-21b5-4732-b95d-619462b4aaf6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = agent.chat(\"Tell me the results over one of the above datasets.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cc4e983",
   "metadata": {},
   "source": [
    "## Lower-Level: Debuggability and Control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55abad72-b189-471a-accc-1621fd19c804",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent_worker = FunctionCallingAgentWorker.from_tools(\n",
    "    [vector_tool, summary_tool], \n",
    "    llm=llm, \n",
    "    verbose=True\n",
    ")\n",
    "agent = AgentRunner(agent_worker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18e911aa-4640-4f89-99c8-6cdf6aff07c6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "task = agent.create_task(\n",
    "    \"Tell me about the agent roles in MetaGPT, \"\n",
    "    \"and then how they communicate with each other.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eaf0b88-e03a-4dd9-91f6-f6f0c8758e64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "step_output = agent.run_step(task.task_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8e77fac-8734-4071-a672-b3a9f30e2bf1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "completed_steps = agent.get_completed_steps(task.task_id)\n",
    "print(f\"Num completed for task {task.task_id}: {len(completed_steps)}\")\n",
    "print(completed_steps[0].output.sources[0].raw_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db8de410-4b82-4daf-93da-28da57cbb0bb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "upcoming_steps = agent.get_upcoming_steps(task.task_id)\n",
    "print(f\"Num upcoming steps for task {task.task_id}: {len(upcoming_steps)}\")\n",
    "upcoming_steps[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc352582-2c17-46ef-ba80-0f571e920c3c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "step_output = agent.run_step(\n",
    "    task.task_id, input=\"What about how agents share information?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be80661f-81b1-45fc-b0ba-33a04dae849b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "step_output = agent.run_step(task.task_id)\n",
    "print(step_output.is_last)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4496328c-e6d5-4722-a8df-78a73a441b3c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = agent.finalize_response(task.task_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "601d1bed-78b2-4512-87ac-aec5ce5d8494",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
